#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A股市场股票查询分析器
功能包括：
1. 查找A股市场中的股票
2. 指定相关概念股查询
3. 指定上市年份查询
4. 指定股价范围查询
5. 指定市值范围查询
6. 查询市盈为正且最近一个月走势向上的股票
"""

import akshare as ak
import pandas as pd
import numpy as np
import warnings
from datetime import datetime, timedelta
from typing import List, Optional, Tuple
import time
import sys

warnings.filterwarnings('ignore')


class StockAnalyzer:
    """A股股票分析器"""
    
    def __init__(self):
        self.stock_info = None
        self.concept_stocks = None
        self.load_basic_data()
    
    def load_basic_data(self):
        """加载基础股票数据"""
        print("正在加载A股基础数据...")
        try:
            # 获取A股基本信息
            self.stock_info = ak.stock_info_a_code_name()
            print(f"成功加载 {len(self.stock_info)} 只A股股票基础信息")
            
            # 获取概念股信息
            print("正在加载概念股信息...")
            try:
                self.concept_stocks = ak.stock_board_concept_name_em()
                print(f"成功加载 {len(self.concept_stocks)} 个概念板块")
            except Exception as e:
                print(f"加载概念股信息失败: {e}")
                self.concept_stocks = pd.DataFrame()
                
        except Exception as e:
            print(f"加载基础数据失败: {e}")
            sys.exit(1)
    
    def get_stock_realtime_data(self, stock_codes: List[str]) -> pd.DataFrame:
        """获取股票实时数据"""
        print(f"正在获取 {len(stock_codes)} 只股票的实时数据...")
        
        all_data = []
        batch_size = 50  # 批量处理以提高效率
        
        for i in range(0, len(stock_codes), batch_size):
            batch_codes = stock_codes[i:i + batch_size]
            print(f"处理第 {i//batch_size + 1} 批，共 {len(batch_codes)} 只股票...")
            
            for code in batch_codes:
                try:
                    # 获取实时数据
                    realtime_data = ak.stock_zh_a_spot_em()
                    stock_data = realtime_data[realtime_data['代码'] == code]
                    
                    if not stock_data.empty:
                        all_data.append(stock_data.iloc[0])
                    
                    time.sleep(0.1)  # 避免请求过于频繁
                    
                except Exception as e:
                    print(f"获取股票 {code} 数据失败: {e}")
                    continue
        
        if all_data:
            return pd.DataFrame(all_data)
        else:
            return pd.DataFrame()
    
    def get_stock_monthly_trend(self, stock_code: str) -> bool:
        """
        判断股票最近一个月走势是否向上
        返回 True 表示向上趋势，False 表示向下趋势
        """
        try:
            # 获取最近60天的日K线数据（确保有足够数据）
            end_date = datetime.now().strftime('%Y%m%d')
            start_date = (datetime.now() - timedelta(days=60)).strftime('%Y%m%d')
            
            # 获取历史数据
            hist_data = ak.stock_zh_a_hist(symbol=stock_code, 
                                         period="daily", 
                                         start_date=start_date, 
                                         end_date=end_date)
            
            if len(hist_data) < 20:  # 数据不足
                return False
            
            # 获取最近20天的收盘价
            recent_prices = hist_data['收盘'].tail(20).values
            
            # 计算线性回归斜率判断趋势
            x = np.arange(len(recent_prices))
            slope = np.polyfit(x, recent_prices, 1)[0]
            
            # 计算价格变化率
            price_change_rate = (recent_prices[-1] - recent_prices[0]) / recent_prices[0]
            
            # 如果斜率为正且价格变化率大于-5%，认为是向上趋势
            return slope > 0 and price_change_rate > -0.05
            
        except Exception as e:
            print(f"获取股票 {stock_code} 走势数据失败: {e}")
            return False
    
    def search_stocks_by_concept(self, concept_name: str) -> List[str]:
        """根据概念查找股票代码"""
        try:
            concept_stocks = ak.stock_board_concept_cons_em(symbol=concept_name)
            return concept_stocks['代码'].tolist()
        except Exception as e:
            print(f"获取概念股 {concept_name} 失败: {e}")
            return []
    
    def get_available_concepts(self) -> List[str]:
        """获取所有可用概念名称"""
        if self.concept_stocks is not None and not self.concept_stocks.empty:
            return self.concept_stocks['板块名称'].tolist()
        return []
    
    def filter_stocks(self, 
                     concept: Optional[str] = None,
                     listing_year: Optional[int] = None,
                     price_range: Optional[Tuple[float, float]] = None,
                     market_value_range: Optional[Tuple[float, float]] = None,
                     positive_pe_and_uptrend: bool = False) -> pd.DataFrame:
        """
        综合筛选股票
        
        参数:
        concept: 概念股名称
        listing_year: 上市年份（查找此年份之后上市的股票）
        price_range: 股价范围 (最低价, 最高价)
        market_value_range: 市值范围 (最小市值亿元, 最大市值亿元)
        positive_pe_and_uptrend: 是否筛选市盈率为正且最近一月走势向上的股票
        """
        
        # 获取初始股票列表
        if concept:
            print(f"正在查找概念股: {concept}")
            stock_codes = self.search_stocks_by_concept(concept)
            if not stock_codes:
                print(f"未找到概念股 '{concept}' 的相关股票")
                return pd.DataFrame()
            print(f"找到 {len(stock_codes)} 只 {concept} 概念股")
        else:
            stock_codes = self.stock_info['code'].tolist()
        
        # 获取实时数据
        try:
            realtime_data = ak.stock_zh_a_spot_em()
        except Exception as e:
            print(f"获取实时数据失败: {e}")
            return pd.DataFrame()
        
        # 筛选股票代码
        if concept:
            realtime_data = realtime_data[realtime_data['代码'].isin(stock_codes)]
        
        if realtime_data.empty:
            print("没有找到符合条件的股票")
            return pd.DataFrame()
        
        filtered_data = realtime_data.copy()
        
        # 按上市年份筛选
        if listing_year:
            print(f"筛选 {listing_year} 年后上市的股票...")
            try:
                # 获取股票基本信息（包括上市日期）
                stock_basic = ak.stock_info_a_code_name()
                # 这里简化处理，实际需要获取详细的上市日期信息
                print("注意：上市年份筛选功能需要更详细的数据接口支持")
            except Exception as e:
                print(f"获取上市日期信息失败: {e}")
        
        # 按股价范围筛选
        if price_range:
            min_price, max_price = price_range
            print(f"筛选股价在 {min_price}-{max_price} 元范围内的股票...")
            filtered_data = filtered_data[
                (filtered_data['最新价'] >= min_price) & 
                (filtered_data['最新价'] <= max_price)
            ]
            print(f"股价筛选后剩余 {len(filtered_data)} 只股票")
        
        # 按市值范围筛选
        if market_value_range:
            min_mv, max_mv = market_value_range
            print(f"筛选市值在 {min_mv}-{max_mv} 亿元范围内的股票...")
            # 市值单位转换（亿元）
            filtered_data = filtered_data[
                (filtered_data['总市值'] >= min_mv * 100000000) & 
                (filtered_data['总市值'] <= max_mv * 100000000)
            ]
            print(f"市值筛选后剩余 {len(filtered_data)} 只股票")
        
        # 筛选市盈率为正且最近一月走势向上的股票
        if positive_pe_and_uptrend:
            print("筛选市盈率为正的股票...")
            # 筛选市盈率为正
            filtered_data = filtered_data[filtered_data['市盈率-动态'] > 0]
            print(f"市盈率筛选后剩余 {len(filtered_data)} 只股票")
            
            if not filtered_data.empty:
                print("分析股票走势趋势...")
                uptrend_stocks = []
                
                for index, row in filtered_data.iterrows():
                    stock_code = row['代码']
                    print(f"分析股票 {stock_code} ({row['名称']}) 的走势...")
                    
                    if self.get_stock_monthly_trend(stock_code):
                        uptrend_stocks.append(index)
                    
                    time.sleep(0.2)  # 避免请求过于频繁
                
                filtered_data = filtered_data.loc[uptrend_stocks]
                print(f"走势筛选后剩余 {len(filtered_data)} 只股票")
        
        return filtered_data
    
    def display_results(self, results: pd.DataFrame):
        """显示筛选结果"""
        if results.empty:
            print("没有找到符合条件的股票")
            return
        
        print(f"\n找到 {len(results)} 只符合条件的股票：")
        print("=" * 100)
        
        # 选择要显示的列
        display_columns = ['代码', '名称', '最新价', '涨跌幅', '市盈率-动态', '总市值']
        
        # 检查列是否存在
        available_columns = [col for col in display_columns if col in results.columns]
        display_data = results[available_columns].copy()
        
        # 市值单位转换为亿元
        if '总市值' in display_data.columns:
            display_data['总市值(亿元)'] = (display_data['总市值'] / 100000000).round(2)
            display_data = display_data.drop('总市值', axis=1)
        
        # 设置显示格式
        pd.set_option('display.max_columns', None)
        pd.set_option('display.width', None)
        pd.set_option('display.max_colwidth', 20)
        
        for index, row in display_data.iterrows():
            print(f"股票代码: {row['代码']}")
            print(f"股票名称: {row['名称']}")
            print(f"最新价格: {row['最新价']:.2f} 元")
            if '涨跌幅' in row:
                print(f"涨跌幅: {row['涨跌幅']:.2f}%")
            if '市盈率-动态' in row:
                print(f"市盈率: {row['市盈率-动态']:.2f}")
            if '总市值(亿元)' in row:
                print(f"总市值: {row['总市值(亿元)']:.2f} 亿元")
            print("-" * 50)


def main():
    """主函数"""
    print("A股市场股票查询分析器")
    print("=" * 50)
    
    # 初始化分析器
    analyzer = StockAnalyzer()
    
    while True:
        print("\n请选择查询功能:")
        print("1. 查看所有A股股票")
        print("2. 按概念股查询")
        print("3. 按股价范围查询")
        print("4. 按市值范围查询")
        print("5. 查询市盈率为正且最近一月走势向上的股票")
        print("6. 综合条件查询")
        print("7. 查看可用概念股列表")
        print("0. 退出")
        
        choice = input("\n请输入选择 (0-7): ").strip()
        
        if choice == '0':
            print("感谢使用，再见！")
            break
        
        elif choice == '1':
            print("正在获取所有A股股票信息...")
            try:
                all_stocks = ak.stock_zh_a_spot_em()
                analyzer.display_results(all_stocks.head(50))  # 显示前50只股票
                print(f"\n注意：共有 {len(all_stocks)} 只股票，此处仅显示前50只")
            except Exception as e:
                print(f"获取股票信息失败: {e}")
        
        elif choice == '2':
            concept_name = input("请输入概念股名称 (例如: 新能源汽车): ").strip()
            if concept_name:
                results = analyzer.filter_stocks(concept=concept_name)
                analyzer.display_results(results)
        
        elif choice == '3':
            try:
                min_price = float(input("请输入最低股价 (元): ").strip())
                max_price = float(input("请输入最高股价 (元): ").strip())
                results = analyzer.filter_stocks(price_range=(min_price, max_price))
                analyzer.display_results(results)
            except ValueError:
                print("请输入有效的数字")
        
        elif choice == '4':
            try:
                min_mv = float(input("请输入最小市值 (亿元): ").strip())
                max_mv = float(input("请输入最大市值 (亿元): ").strip())
                results = analyzer.filter_stocks(market_value_range=(min_mv, max_mv))
                analyzer.display_results(results)
            except ValueError:
                print("请输入有效的数字")
        
        elif choice == '5':
            print("正在查询市盈率为正且最近一月走势向上的股票...")
            print("这可能需要较长时间，请耐心等待...")
            results = analyzer.filter_stocks(positive_pe_and_uptrend=True)
            analyzer.display_results(results)
        
        elif choice == '6':
            print("综合条件查询")
            
            # 概念股
            concept = input("概念股名称 (可选，直接回车跳过): ").strip()
            concept = concept if concept else None
            
            # 上市年份
            listing_year_input = input("上市年份之后 (可选，直接回车跳过): ").strip()
            listing_year = int(listing_year_input) if listing_year_input else None
            
            # 股价范围
            price_range = None
            price_input = input("股价范围 (格式: 最低价,最高价，可选): ").strip()
            if price_input:
                try:
                    prices = price_input.split(',')
                    price_range = (float(prices[0]), float(prices[1]))
                except:
                    print("股价格式错误，跳过此条件")
            
            # 市值范围
            market_value_range = None
            mv_input = input("市值范围 (格式: 最小市值亿元,最大市值亿元，可选): ").strip()
            if mv_input:
                try:
                    mvs = mv_input.split(',')
                    market_value_range = (float(mvs[0]), float(mvs[1]))
                except:
                    print("市值格式错误，跳过此条件")
            
            # 市盈率和走势
            pe_trend = input("是否筛选市盈率为正且走势向上的股票？ (y/n): ").strip().lower() == 'y'
            
            # 执行查询
            results = analyzer.filter_stocks(
                concept=concept,
                listing_year=listing_year,
                price_range=price_range,
                market_value_range=market_value_range,
                positive_pe_and_uptrend=pe_trend
            )
            analyzer.display_results(results)
        
        elif choice == '7':
            concepts = analyzer.get_available_concepts()
            if concepts:
                print(f"\n可用概念股列表 (共 {len(concepts)} 个):")
                for i, concept in enumerate(concepts[:50], 1):  # 显示前50个
                    print(f"{i:2d}. {concept}")
                if len(concepts) > 50:
                    print(f"... 还有 {len(concepts) - 50} 个概念股")
            else:
                print("无法获取概念股列表")
        
        else:
            print("无效选择，请重新输入")


if __name__ == "__main__":
    main() 